# PiCANet/train.py
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms
from picanet_resnet18 import PiCANetResnet18
from dataset import ECSSDDataset
from tqdm import tqdm

def train(args,device):
    image_dir = os.path.join(args.data_path, 'images')
    mask_dir = os.path.join(args.data_path, 'ground_truth_mask')

    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])
    mask_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])

    dataset = ECSSDDataset(image_dir, mask_dir, transform, mask_transform)
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)

    model = PiCANetResnet18().to(device)

    criterion = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

    os.makedirs(args.save_path, exist_ok=True)

    for epoch in range(args.epochs):
        model.train()
        running_loss = 0.0
        loop = tqdm(dataloader, desc=f"Epoch {epoch+1}/{args.epochs}")
        for images, masks in loop:
            images = images.to(device)
            masks = masks.to(device)

            x1, x2, x3, x4 = model(images)

            pred = nn.Conv2d(512, 1, kernel_size=1).to(device)(x4)
            pred = nn.functional.interpolate(pred, size=(224, 224), mode='bilinear', align_corners=False)

            loss = criterion(pred, masks)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            loop.set_postfix(loss=loss.item())

        print(f"Epoch [{epoch+1}/{args.epochs}], Loss: {running_loss / len(dataloader):.4f}")

        torch.save(model.state_dict(), os.path.join(args.save_path, f"picanet_epoch{epoch+1}.pth"))
